/******************************************************************************
 This is a derivative work of the data structure presented in:
 
 http://www.cmcrossroads.com/bradapp/ftp/src/libs/C++/AvlTrees.html

 whose author is Brad Appleton (http://www.bradapp.net/). 
 
 See License.txt for license information about the use of the data
 structure.

 Please, note the links above were valid when this comment was written and 
 may not be valid at the time you read it.
 ******************************************************************************/

#include <stdlib.h>
#include <algorithm>

#include "avl_iterator.h"

// ---------------------------------------------------------------- Definitions

namespace cnome { namespace tst_avl {
// Return the minumum of two numbers
template<typename T> T
    MIN(T a, T b) {
    return  (a < b) ? a : b;
}

// Return the maximum of two numbers
template<typename T> T
    MAX(T a, T b) {
    return  (a > b) ? a : b;
}

// Use mnemonic constants for valid balance-factor values
enum balance_t { LEFT_HEAVY = -1, BALANCED = 0, RIGHT_HEAVY = 1 };

// Use mnemonic constants for indicating a change in height
enum height_effect_t { HEIGHT_NOCHANGE = 0, HEIGHT_CHANGE = 1 };

// Return true if the tree is too heavy on the left side
inline static int
    LEFT_IMBALANCE(short bal) {
    return (bal < LEFT_HEAVY);
}

// Return true if the tree is too heavy on the right side
inline static int
    RIGHT_IMBALANCE(short bal) {
    return (bal > RIGHT_HEAVY);
}

// ----------------------------------------------- Constructors and Destructors
template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>::avlNode(std::size_t keySize)
    : nodeInfo(), parent(0), left(0), middle(0), right(0), data(0) {
    _ASSERT(keySize>=0);
    if (keySize == 0) {
        nodeInfo |= NILLMASK;
        parent = this;
        subtree(LEFT) = this;
//            subtree(MIDDLE) = this;
        subtree(RIGHT) = this;
    }
    else {
        if (keySize > 1) {
#ifdef max
#undef max
#endif
            size_t lengthOfCharsLength;
            if (keySize <= std::numeric_limits<unsigned char>::max())
                lengthOfCharsLength = 1;
            else if (keySize <  std::numeric_limits<unsigned short>::max())
                lengthOfCharsLength = 2;
            else if (keySize <  std::numeric_limits<unsigned long>::max())
                lengthOfCharsLength = 3;
            else
                throw std::invalid_argument("Invalid argument value");

            nodeInfo |= (lengthOfCharsLength << 3);
        }

        setKeySize(keySize);
    }
#ifdef AVL_DEBUG
    keyDBG = keyAddress();
    balDBG = 0;
#endif
}

template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>::avlNode(const avlNode<Char, T, Pred, Alloc> &a)
    : nodeInfo(a.nodeInfo), parent(a.parent), left(a.left), middle(a.middle), right(a.right), data(a.data) {
    setKey(a.getKey());
#ifdef AVL_DEBUG
    keyDBG = keyAddress();
    keySzDBG = 1;
    balDBG = a.balDBG;
#endif
}

template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>::avlNode(const avlNode<Char, T, Pred, Alloc> &a, std::size_t keySize)
    : nodeInfo(a.nodeInfo&avlNode<Char, T, Pred, Alloc>::BALANCEMASK),
      left(a.left), middle(a.middle), right(a.right), data(a.data) {
    setKeySize(MIN(a.keySize(), keySize));
    memcpy(keyAddress(), a.keyAddress(), keySize);

#ifdef AVL_DEBUG
    keyDBG = keyAddress();
    balDBG = a.balDBG;
#endif
}

template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>::~avlNode(void) {
    operator delete(subtree(LEFT), 0);
    operator delete(subtree(MIDDLE), 0);
    operator delete(subtree(RIGHT), 0);
}

// ------------------------------------------------- Rotating and Re-Balancing

template <typename Char, typename T, typename Pred, typename Alloc>
int
    avlNode<Char, T, Pred, Alloc>::rotateOnce(avlNode<Char, T, Pred, Alloc> * & root, dir_t dir) {
    dir_t  otherDir = opposite(dir);
    avlNode * oldRoot = root;

    // See if otherDir subtree is balanced. If it is, then this
    // rotation will *not* change the overall tree height.
    // Otherwise, this rotation will shorten the tree height.
    int heightChange =   (root->subtree(otherDir)->getBal() == 0)
                       ? HEIGHT_NOCHANGE
                       : HEIGHT_CHANGE;

    // assign new root
    root = oldRoot->subtree(otherDir);
    root->parent = oldRoot->parent;

    // new-root exchanges it's "dir" mySubtree for it's parent
    oldRoot->subtree(otherDir) = root->subtree(dir);
    if (root->subtree(dir))
        root->subtree(dir)->parent = oldRoot;
    root->subtree(dir) = oldRoot;
    oldRoot->parent = root;

    // update balances
    oldRoot->setBal(-((dir == LEFT) ? root->decBal() : root->incBal()));

    return  heightChange;
}

template <typename Char, typename T, typename Pred, typename Alloc>
int
    avlNode<Char, T, Pred, Alloc>::rotateTwice(avlNode<Char, T, Pred, Alloc> * & root, dir_t dir) {
    dir_t  otherDir = opposite(dir);
    avlNode* oldRoot = root;
    avlNode* oldOtherDirSubtree = root->subtree(otherDir);

    // assign new root
    avlNode* lParent = root->parent;
    root = oldRoot->subtree(otherDir)->subtree(dir);
    root->parent = lParent;

    // new-root exchanges it's "dir" mySubtree for it's grandparent
    oldRoot->subtree(otherDir) = root->subtree(dir);
    if (root->subtree(dir))
        root->subtree(dir)->parent = oldRoot;

    root->subtree(dir) = oldRoot;
    oldRoot->parent = root;

    // new-root exchanges it's "other-dir" mySubtree for it's parent
    oldOtherDirSubtree->subtree(dir) = root->subtree(otherDir);
    if (root->subtree(otherDir))
        root->subtree(otherDir)->parent = oldOtherDirSubtree;
    root->subtree(otherDir) = oldOtherDirSubtree;
    oldOtherDirSubtree->parent = root;

    // update balances
    root->subtree(LEFT)->setBal(-MAX(root->getBal(), static_cast<signed char>(0)));
    root->subtree(RIGHT)->setBal(-MIN(root->getBal(),static_cast<signed char>(0)));
    root->setBal(0);

    // A double rotation always shortens the overall height of the tree
    return  HEIGHT_CHANGE;
}

template <typename Char, typename T, typename Pred, typename Alloc>
int
    avlNode<Char, T, Pred, Alloc>::reBalance(avlNode<Char, T, Pred, Alloc> * & root) {
    int  heightChange = HEIGHT_NOCHANGE;

    if (LEFT_IMBALANCE(root->getBal())) {   // Need a right rotation
        if (root->subtree(LEFT)->getBal() == RIGHT_HEAVY)
            heightChange = rotateTwice(root, RIGHT); // RL rotation needed
        else 
            heightChange = rotateOnce(root, RIGHT);  // RR rotation needed
    } 
    else if (RIGHT_IMBALANCE(root->getBal())) {   // Need a left rotation
        if (root->subtree(RIGHT)->getBal() == LEFT_HEAVY) 
            heightChange = rotateTwice(root, LEFT);  // LR rotation needed
        else 
            heightChange = rotateOnce(root, LEFT);   // LL rotation needed
    }

    return  heightChange;
}
// ------------------------------------------------------- Comparisons

template <typename Char, typename T, typename Pred, typename Alloc>
cmp_t
    avlNode<Char, T, Pred, Alloc>::compare(const Pred& pred, Char key, cmp_t cmp) const {
    switch (cmp) {
        case EQ_CMP :  // Standard comparison
            return    pred(key, getKey()) ? MIN_CMP
                    : pred(getKey(), key) ? MAX_CMP
                    : EQ_CMP;

        case MIN_CMP :  // Find the minimal element in this tree
            return  !subtree(LEFT) ? EQ_CMP : MIN_CMP;

        case MAX_CMP :  // Find the maximal element in this tree
            return  !subtree(RIGHT) ? EQ_CMP : MAX_CMP;
    }

    return EQ_CMP;
}

// ------------------------------------------------------- search/insert/erase

template <typename Char, typename T, typename Pred, typename Alloc>
T *
    avlNode<Char, T, Pred, Alloc>::search(const Pred& pred, Char key, avlNode<Char, T, Pred, Alloc>* root, cmp_t cmp) {
    cmp_t result;
    while (root  &&  (result = root->compare(pred, key, cmp))) {
        root = root->subtree((result < 0) ? LEFT : RIGHT);
    }
    return  (root) ? root->getData() : 0;
}

template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>*& 
    avlNode<Char, T, Pred, Alloc>::find(const Pred& pred, Char key, const avlNode<Char, T, Pred, Alloc>*& root, cmp_t cmp) {
    avlNode** node = const_cast<avlNode**>(&root);
    cmp_t result;
    while (*node  &&  (result = (*node)->compare(pred, key, cmp))) {
        node = &(*node)->subtree((result < 0) ? LEFT : RIGHT);
    }

    return *node;
}

template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>*& 
    avlNode<Char, T, Pred, Alloc>::findEx(Pred pred, Char key, avlNode<Char, T, Pred, Alloc>*& root,
                                          avlNode<Char, T, Pred, Alloc>**& prtRoot, cmp_t cmp) {
    prtRoot = &(root?root->parent->myLinkInParent():root);
    avlNode** node = &root;
    cmp_t result = MIN_CMP;
    while (*node  &&  (result = (*node)->compare(pred, key, cmp))) {
        prtRoot = node;
        node = &(*node)->subtree((result < 0) ? LEFT : RIGHT);
    }

    return *node;
}

template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>*& 
    avlNode<Char, T, Pred, Alloc>::insert(const Pred& pred, Char* key, std::size_t keySize, 
    T *item, avlNode<Char, T, Pred, Alloc>*& root) {
    avlNode** prtNode;
    avlNode** node = &findEx(pred, *key, root, prtNode);

    if (*node) // Found?
        return *node;

    // Not found

    // insert new node here 
    avlNode* newNode = avlNode::createAvlNode(keySize);
    newNode->setKey(key, keySize); 
    newNode->getData() = item;

    if (!root)
        return root = newNode;

    newNode->parent = *prtNode;
    *node = newNode;

    int change =  HEIGHT_CHANGE;
    node = prtNode;
    //while (node && *node && change)
    while (!(*node)->isNill() && change) {
        prtNode = &(*prtNode)->parent->myLinkInParent();
        cmp_t result = (*node)->compare(pred, *key);
        dir_t dir = result==MIN_CMP ? LEFT : RIGHT;
        int increase = (dir == LEFT?MIN_CMP:MAX_CMP) * change;
        (*node)->setBal((*node)->getBal() + increase);       // update balance factor 

        // ----------------------------------------------------------------------
        // re-balance if needed -- height of current tree increases only if its
        // subtree height increases and the current tree needs no rotation.
        // ----------------------------------------------------------------------
        change =   (increase && (*node)->getBal())
                 ? (1 - reBalance(*node))
                 : HEIGHT_NOCHANGE;

        if ((*node)->myDirInParent() == MIDDLE)
            break;

        node = prtNode;
    }
    return newNode->myLinkInParent();
}

template <typename Char, typename T, typename Pred, typename Alloc>
T* 
avlNode<Char, T, Pred, Alloc>::erase(const Pred& pred, avlNode<Char, T, Pred, Alloc>*& root) {
    if (!root)
        return 0;

    avlNode** node = &root;
    avlNode** prtNode = &root->parent->myLinkInParent();
    dir_t dirP = (*node)->myDirInParent();

    _ASSERT((*prtNode)->subtree((*node)->myDirInParent()) == *node);

    // Found key at this node
    T * found = (*node)->getData();
    int  decrease = 0;

    // ---------------------------------------------------------------------
    // At this point we know "result" is zero and "root" points to
    // the node that we need to delete.  There are three cases:
    //
    //    1) The node is a leaf.  Remove it and return.
    //
    //    2) The node is a branch (has only 1 child). Make "root"
    //       (the pointer to this node) point to the child.
    //
    //    3) The node has two children. We swap items with the successor
    //       of "root" (the smallest item in its right subtree) and delete
    //       the successor from the right subtree of "root".  The
    //       identifier "decrease" should be reset if the subtree height
    //       decreased due to the deletion of the successor of "root".
    // ---------------------------------------------------------------------

    if (    !(*node)->subtree(LEFT)  
         && !(*node)->subtree(RIGHT) ) {
        // We have a leaf -- remove it
        decrease = ((*node)->myDirInParent() == RIGHT)? HEIGHT_CHANGE : -HEIGHT_CHANGE;    // height changed from 1 to 0
        (*node)->parent = 0;

        operator delete((*node), 0);
        *node = 0;
    } 
    else if (    !(*node)->subtree(LEFT) 
              || !(*node)->subtree(RIGHT)) {
        // We have one child -- only child becomes new root 
        avlNode* toDelete = (*node);
        *node = (*node)->subtree(((*node)->subtree(RIGHT)) ? RIGHT : LEFT);
        (*node)->parent = *prtNode;
        decrease = ((*node)->myDirInParent() == RIGHT)? HEIGHT_CHANGE : -HEIGHT_CHANGE;    // We just shortened the subtree
        // Null-out the subtree pointers so we don't recursively delete
        toDelete->parent = 0;
        toDelete->subtree(LEFT) = 0;
        toDelete->subtree(RIGHT) = 0;

        operator delete(toDelete, 0);
    }
    else {
		// We have two children -- find successor and replace our current
        // data item with that of the successor
        //DataEnsured(*node) = erase(key, (*node)->subtree(RIGHT), decrease, MIN_CMP);
        // TODO: Remove recursive call above
        avlNode* toDelete = *node;
        avlNode* toDeleteL = (*node)->subtree(LEFT);
        avlNode* toDeleteR = (*node)->subtree(RIGHT);
			avlNode** node2 = &find(pred, (*node)->getKey(), const_cast<const avlNode*&>((*node)->subtree(RIGHT)), MIN_CMP);
			avlNode* findToReplace = *node2;
        _ASSERT(!findToReplace->subtree(LEFT));

        decrease = (findToReplace->myDirInParent() == RIGHT)? HEIGHT_CHANGE : -HEIGHT_CHANGE;

			// Remove findToReplace from its position
			if (findToReplace->subtree(RIGHT))
				findToReplace->subtree(RIGHT)->parent = findToReplace->parent;
			*node2 = findToReplace->subtree(RIGHT);

			// Repos findToReplace in the toDelete pos
			findToReplace->parent = toDelete->parent;
			*node = findToReplace;
			findToReplace->subtree(RIGHT) = toDelete->subtree(RIGHT);
			if (toDelete->subtree(RIGHT))
				toDelete->subtree(RIGHT)->parent = findToReplace;
			findToReplace->subtree(LEFT) = toDelete->subtree(LEFT);
			toDelete->subtree(LEFT)->parent = findToReplace;
			findToReplace->setBal(toDelete->getBal());

        // Null-out the subtree pointers so we don't recursively delete
        toDelete->parent = 0;
        toDelete->subtree(LEFT) = 0;
        toDelete->subtree(RIGHT) = 0;

			// Assign prtNode
			prtNode = &(*node2?(*node2)->parent:findToReplace)->myLinkInParent();

        operator delete(toDelete, 0);
    }

    node = prtNode;
    while (!(*node)->isNill() && decrease && dirP != MIDDLE) {
        prtNode = &(*prtNode)->parent->myLinkInParent();
        (*node)->setBal((*node)->getBal() - decrease);       // update balance factor 

        // ------------------------------------------------------------------------
        // Rebalance if necessary -- the height of current tree changes if one
        // of two things happens: (1) a rotation was performed which changed
        // the height of the subtree (2) the subtree height decreased and now
        // matches the height of its other subtree (so the current tree now
        // has a zero balance when it previously did not).
        // ------------------------------------------------------------------------
        //decrease = (decrease) ? ((root->myBal) ? balance(root) : HEIGHT_CHANGE)
        //                      : HEIGHT_NOCHANGE ;
        if ((*node)->getBal())           // rebalance and see if height changed
            decrease = ((*node)->myDirInParent() == RIGHT)? reBalance(*node) : -reBalance(*node);  
        else                                     // balanced because subtree decreased
            decrease = ((*node)->myDirInParent() == RIGHT)? HEIGHT_CHANGE : -HEIGHT_CHANGE;   
        
        if ((*prtNode)->isNill())
            break;

        dirP = (*node)->myDirInParent();
        node = prtNode;
    }

    return found;
}

template <typename Char, typename T, typename Pred, typename Alloc>
T *
avlNode<Char, T, Pred, Alloc>::erase(const Pred& pred, Char key, avlNode<Char, T, Pred, Alloc>*& root,        
                                      cmp_t cmp) {
    return erase(pred, find(pred, key, root, cmp));
}

// --------------------------------------------------------------- Verification

template <typename Char, typename T, typename Pred, typename Alloc>
int
    avlNode<Char, T, Pred, Alloc>::height() const {
    int  leftHeight  = (subtree(LEFT)) ? subtree(LEFT)->height()  : 0;
    int  rightHeight = (subtree(RIGHT)) ? subtree(RIGHT)->height() : 0;
    return  (1 + MAX(leftHeight, rightHeight));
}

template <typename Char, typename T, typename Pred, typename Alloc>
int
    avlNode<Char, T, Pred, Alloc>::check(const Pred& pred) const {
    int  valid = 1;

    // First verify that subtrees are correct
    if (subtree(LEFT))  valid *= subtree(LEFT)->check(pred);
    if (subtree(RIGHT)) valid *= subtree(RIGHT)->check(pred);

    // Now get the height of each subtree
    int  leftHeight  = subtree(LEFT)  ? subtree(LEFT)->height()   : 0;
    int  rightHeight = subtree(RIGHT) ? subtree(RIGHT)->height()  : 0;

    // Verify that AVL tree property is satisfied
    int  diffHeight = rightHeight - leftHeight;
    if (LEFT_IMBALANCE(diffHeight) || RIGHT_IMBALANCE(diffHeight)) {
        valid = 0;
        std::cerr << "height difference is " << diffHeight
            << " at node " << getKey() << std::endl;
    }

    // Verify that balance-factor is correct
    if (diffHeight != getBal()) {
        valid = 0;
        std::cerr << "height difference " << diffHeight
            << " doesn't match balance-factor of " << static_cast<short>(getBal())
            << " at node " << getKey() << std::endl;
    }

    // Verify that search-tree property is satisfied
    if (    subtree(LEFT) 
        && subtree(LEFT)->compare(pred, getKey()) == MIN_CMP) {
        valid = 0;
        std::cerr << "Node " << getKey() << " is *smaller* than left subtree "
            << subtree(LEFT)->getKey() << std::endl;
    }
    if (    subtree(RIGHT)
        && subtree(RIGHT)->compare(pred, getKey()) == MAX_CMP) {
        valid = 0;
        std::cerr << "Node " << getKey() << " is *greater* than right subtree "
            << subtree(RIGHT)->getKey() << std::endl;
    }

    // Verify the parent link
    if (parent)
        if (this != parent->subtree(LEFT) && this != parent->subtree(MIDDLE) && this != parent->subtree(RIGHT)) {
            valid = 0;
            std::cerr << "Node " << getKey() << " is not child of its father " << parent->getKey() << std::endl;
        }

    return  valid;
}

template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>* avlNode<Char, T, Pred, Alloc>::shrinkNode(avlNode*& node) {
    return node;
}

//-- operators
template <typename Char, typename T, typename Pred, typename Alloc>
avlNode<Char, T, Pred, Alloc>&  
    avlNode<Char, T, Pred, Alloc>::operator=(const avlNode<Char, T, Pred, Alloc> & a) {
    unsigned char lengthOfCharsLength(getLengthOfCharsLength());
    unsigned char a_lengthOfCharsLength(a.getLengthOfCharsLength());

    if (keySize() != a.keySize())
        throw std::invalid_argument("Key size mismatch");  // Use copy constructor

    setKey(a.keyAddress(), keySize());

    parent = a.parent;
    subtree(LEFT) = a.subtree(LEFT);
    subtree(MIDDLE) = a.subtree(MIDDLE);
    subtree(RIGHT) = a.subtree(RIGHT);
    getData() = a.getData();
    nodeInfo = (nodeInfo & 0xFC) | (a.nodeInfo & 0x03); // TODO: Verify this command

    return *this;
}

template <typename Char, typename T, typename Pred, typename Alloc>
void*  
    avlNode<Char, T, Pred, Alloc>::operator new(size_t s, size_t* s2) {
    size_t finalSz = s + *s2 - sizeof(Char);
    return malloc(finalSz);
}

template <typename Char, typename T, typename Pred, typename Alloc>
void  
    avlNode<Char, T, Pred, Alloc>::operator delete(void * p, size_t*) {
    free(p);
}

//----------------------------------------------- Routines for dumping the tree

static inline std::ostream& Indent(std::ostream & os, int len) {
    for (int i = 0; i < len; i++) {
        os << ' ';
    }
    return  os;
}

enum TraversalOrder { LTREE, KEY, RTREE };

template <typename Char, typename T, typename Pred, typename Alloc>
static void
    Dump(std::ostream & os,
    TraversalOrder order,
    const avlNode<Char, T, Pred, Alloc> * node,
    int level=0) {
    unsigned  len = (level * 5) + 1;
    if ((order == LTREE) && (!node->subtree(LEFT))) {
        Indent(os, len) << "     **NULL**" << std::endl;
    }
    if (order == KEY) {
        Indent(os, len) << node->getKey() << ":" << static_cast<short>(node->getBal()) << std::endl;
    }
    if ((order == RTREE) && (!node->subtree(RIGHT))) {
        Indent(os, len) << "     **NULL**" << std::endl;
    }
}

template <typename Char, typename T, typename Pred, typename Alloc>
static void
    Dump(std::ostream & os, const avlNode<Char, T, Pred, Alloc> * node, int level=0) {
    if (!node || node->isNill()) {
        os << "***EMPTY TREE***" << std::endl;
    } 
    else {
        Dump(os, RTREE, node, level);
        if (node->subtree(RIGHT)) {
            Dump(os, node->subtree(RIGHT), level+1);
        }
        Dump(os, KEY, node, level);
        if (node->subtree(LEFT)) {
            Dump(os, node->subtree(LEFT), level+1);
        }
        Dump(os, LTREE, node, level);
    }// if non-empty tree
}

template <typename Char, typename T, typename Pred, typename Alloc>
void
    avlTree<Char, T, Pred, Alloc>::dumpTree(std::ostream & os) const {
    Dump(os, treeRoot());
}
}}  // cnome::tst_avl
